import numpy as np
import math
import matplotlib.pyplot as plt


"""
    notes: all coordinates are using right-hand Descartes System with order XYZ
"""


def lineLineIntersect(p0: np.ndarray, d0: np.ndarray, p1: np.ndarray, d1: np.ndarray) -> np.ndarray:
    """
    calculate the intersection point of two lines in 2D space

    :param p0: point on line0
    :param d0: direction vector of line0
    :param p1: point on line1
    :param d1: direction vector of line1
    :return: intersection point, (2,) ndarray
    """
    # TODO: used to calculate from graph draft
    # n0 = np.asarray([-d0[1], d0[0]])
    # if n0.dot(d1) > 0:
    #     n0 *= -1
    # elif n0.dot(d1) == 0:
    #     print("something wrong in lineLineIntersect, line parallel")
    # # x = (p1 - p0).dot(n0) / math.sin(math.acos(d1.dot(d0)))
    # x = (p1 - p0).dot(n0) / math.sqrt(1 - d1.dot(d0) ** 2)
    # x *= d1
    # x += p1
    # return x

    x1, y1 = p0
    x2, y2 = p1
    m1, n1 = d0
    m2, n2 = d1

    t = ((y1 - y2) * m2 - (x1 - x2) * n2) / (m1 * n2 - m2 * n1)
    x = x1 + t * m1
    y = y1 + t * n1
    return np.asarray([x, y])


def linePlaneIntersect(p0: np.ndarray, d0: np.ndarray, p1: np.ndarray, n1: np.ndarray) -> np.ndarray:
    """
    calculate the intersection point of one line and one plane in 3D space

    :param p0: point on line
    :param d0: direction of line
    :param p1: point on plane
    :param n1: normal of plane
    :return: intersection point, (3,) ndarray
    """
    x = (p1 - p0).dot(n1) / d0.dot(n1)
    x *= d0
    x += p0
    return x


def planePlaneIntersect(p0: np.ndarray, n0: np.ndarray, p1: np.ndarray, n1: np.ndarray) -> tuple:
    """
    calculate the intersection line of two planes in 3D space

    :param p0:
    :param n0:
    :param p1:
    :param n1:
    :return: point, direction
    """
    target_d = np.cross(n0, n1)
    target_d /= np.linalg.norm(target_d)
    d0 = np.cross(n0, target_d)
    d0 /= np.linalg.norm(d0)
    target_p = linePlaneIntersect(p0, d0, p1, n1)
    return target_p, target_d


def get_rotation_angle(vector: np.ndarray):
    """
    get rotation angle (alpha, beta) to rotate z-axis to the vector direction
    in form of YZ-EulerAngle using inertial coordinate system.

    :param vector: target direction
    :return: alpha(relative to y-axis), beta(relative to z-axis)
    """

    # use the projection on XZ plane
    # src vector: z-axis, target vector: vector_xz

    # vector_xz = vector.copy()
    # vector_xz[1] = 0
    # z_hat = np.asarray([0, 0, 1])
    # alpha = z_hat.dot(vector_xz) / np.linalg.norm(vector_xz)
    # alpha = math.acos(alpha)
    # if np.cross(z_hat, vector_xz)[1] < 0:
    #     alpha *= -1

    # rotate to the proper direction $temp around Y-axis, s.t. $temp and $vector have the same z-component
    # to make sure $temp can be rotated to $vector around Z-axis
    # alpha is always [0,PI]

    alpha = math.acos(vector[2])


    # now, src vector is in the position of vector_xz

    # use the projection on XY plane
    # src vector: vector_xz_xy, target vector: vector_xy

    # vector_xy = vector.copy()
    # vector_xy[2] = 0
    # vector_xz_xy = vector_xz
    # vector_xz_xy[2] = 0
    # beta = vector_xz_xy.dot(vector_xy) / (np.linalg.norm(vector_xz_xy) * np.linalg.norm(vector_xy))
    # beta = math.acos(beta)
    # if np.cross(vector_xz_xy, vector_xy)[2] < 0:
    #     beta *= -1

    # rotate $temp to $vector around Z-axis. 
    # the radian $beta is decided by the XY-component of $temp and $vector

    source = np.zeros(3)
    source[2] = vector[2]
    source[0] = math.sqrt(1 - source[2] ** 2)
    source_xy = source.copy()
    source_xy[2] = 0
    source_xy /= np.linalg.norm(source_xy)

    vector_xy = vector.copy()
    vector_xy[2] = 0
    vector_xy /= np.linalg.norm(vector_xy)

    beta = source_xy.dot(vector_xy)
    beta = math.acos(beta)
    if np.cross(source_xy, vector_xy)[2] < 0:
        beta *= -1

    return alpha, beta


def get_line_distance(p0: np.ndarray, v0: np.ndarray, p1: np.ndarray, v1: np.ndarray):
    """
    calculate the distance between two non-coplanar lines. if the lines are coplanar, return 0.

    :param p0:
    :param v0:
    :param p1:
    :param v1:
    :return: distance
    """
    vertical_dir = np.cross(v0, v1)
    if (vertical_dir == np.zeros(3)).all():
        return 0
    vertical_dir /= np.linalg.norm(vertical_dir)
    p0p1 = p0 - p1
    distance = abs(p0p1.dot(vertical_dir))
    return distance


def get_rodrigues_rotation_matrix(source_t: np.ndarray, target_t: np.ndarray) -> np.ndarray:
    """
    get the 3X3 rotation matrix to rotate source_t to target_t

    :param source_t:
    :param target_t:
    :return:
    """
    s = source_t.copy()
    t = target_t.copy()
    s /= np.linalg.norm(s)
    t /= np.linalg.norm(t)

    k = np.cross(s, t)
    k /= np.linalg.norm(k)
    co_theta = np.dot(s, t)
    Mk = np.asarray([
        [0., -k[2], k[1]],
        [k[2], 0., -k[0]],
        [-k[1], k[0], 0.]
    ])
    R = np.eye(3) * co_theta + (1 - co_theta) * np.outer(k, k) + np.sqrt(1 - co_theta**2) * Mk
    return R


def get_axis_center(p: np.ndarray, d: np.ndarray, l, u):
    """
    calculate the point on the long axis that is in the middle of the heart

    :param p: point of the line
    :param d: direction of the line
    :param l: lower bound of the AABB
    :param u: upper bound of the AABB
    :return:
    """
    t_in = [0, 0, 0]
    t_out = [0, 0, 0]

    for i in range(3):
        if d[i] == 0:
            if p[i] > l[i] and p[i] < u[i]:
                t_in[i] = -float('inf')
                t_out[i] = float('inf')
            else:
                t_in[i] = float('inf')
                t_out[i] = -float('inf')
            continue
        t_in[i] = (l[i] - p[i]) / d[i]
        t_out[i] = (u[i] - p[i]) / d[i]
        if t_in[i] > t_out[i]:
            t_in[i], t_out[i] = t_out[i], t_in[i]
    t1 = max(t_in)
    t2 = min(t_out)

    # print(t1)
    # print(t2)
    assert t1 < t2
    return p + (t1 + t2) / 2 * d


def get_line_function_2D(point1: np.ndarray, point2: np.ndarray):
    """
    return the tuple (A, B, C) referring to the line function Ax+By+C=0, and (A, B) is the positive normal vector of directional line p1->p2

    :param point1: need be np.float32, in XY order
    :param point2:
    :return:
    """
    v = point2 - point1
    R = np.asarray([
        [np.cos(np.pi / 2), -np.sin(np.pi / 2)],
        [np.sin(np.pi / 2), np.cos(np.pi / 2)]
    ])
    n = np.matmul(R, v)
    a, b = n
    c = -a * point1[0] - b * point1[1]
    return a, b, c


if __name__ == '__main__':
    fig, ax = plt.subplots()

    # TODO: test line line intersect
    # point0 = np.asarray([221., 225.30769231])
    # direction0 = np.asarray([-0.80171518, 0.57374282])
    # direction0 /= np.linalg.norm(direction0)
    #
    # point1 = np.asarray([221., 358.96923077])
    # direction1 = np.asarray([0.95891978, 0.14685303])
    # direction1 /= np.linalg.norm(direction1)
    #
    # intersect = lineLineIntersect(point0, direction0, point1, direction1)
    #
    # ex_point0 = point0 + 200 * direction0
    # ex_point1 = point1 + 200 * -direction1
    # ax.add_line(plt.Line2D([point0[0], ex_point0[0]], [point0[1], ex_point0[1]]))
    # ax.add_line(plt.Line2D([point1[0], ex_point1[0]], [point1[1], ex_point1[1]]))
    # ax.scatter([point0[0], point1[0], intersect[0]], [point0[1], point1[1], intersect[1]])

    # TODO: test line plane intersect
    # fig = plt.figure()
    # ax = fig.add_subplot(projection='3d')
    # p0 = np.asarray([0., 0., 0.])
    # d0 = np.asarray([1., 1., 1.])
    # d0 /= np.linalg.norm(d0)
    # p1 = np.asarray([2., 2., 2.])
    # n1 = np.asarray([1., 1., 1.])
    # n1 /= np.linalg.norm(n1)
    #
    # ax.plot_trisurf(np.asarray([0, 0, 6]), np.asarray([0, 6, 0]), np.asarray([6, 0, 0]))
    # ax.plot([0, 3], [0, 3], [0, 3])
    # x = linePlaneIntersect(p0, d0, p1, n1)
    # ax.scatter([x[0]], [x[1]], [x[2]], c='red')

    # TODO: test plane plane intersect
    # fig = plt.figure()
    # ax = fig.add_subplot(projection='3d')
    # p0 = np.asarray([0., 0., 2.])
    # n0 = np.asarray([0., 0., 1.])
    # p1 = np.asarray([2., 2., 2.])
    # n1 = np.asarray([1., 1., 1.])
    # n1 /= np.linalg.norm(n1)
    #
    # ax.plot_trisurf(np.asarray([0, 2, 0, 2]), np.asarray([0, 0, 2, 2]), np.asarray([2, 2, 2, 2]))
    # ax.plot_trisurf(np.asarray([0, 0, 6]), np.asarray([0, 6, 0]), np.asarray([6, 0, 0]))
    # p, d = planePlaneIntersect(p0, n0, p1, n1)
    # pline_ex = p + 5 * d
    # ax.plot([p[0], pline_ex[0]], [p[1], pline_ex[1]], [p[2], pline_ex[2]])

    plt.show()

    # TODO: test the rotation angle
    # target_vector = np.asarray([1.0, 1.0, 1.0])
    # target_vector /= np.linalg.norm(target_vector)
    # a, b = get_rotation_angle(target_vector)
    # print(a, b)

    # TODO: test the rodrigues rotation
    # s = np.asarray([1., 0., 0.])
    # t = np.asarray([1., 1., 1.])
    # print(get_rodrigues_rotation_matrix(s, t))

    # TODO: test the axis center
    p = np.asarray([0., 0., 0.])
    d = np.asarray([1., 0., 0.])
    l = [1., -1., -1.]
    u = [2., 1., 1.]
    print(get_axis_center(p, d, l, u))

    # TODO: test the non-coplanar line distance
    # p0 = np.asarray([0., 0., 0.])
    # v0 = np.asarray([1., 0., 0.])
    # p1 = np.asarray([0., 0., 1.])
    # v1 = np.asarray([0., 1., 1.])
    # v1 /= np.linalg.norm(v1)
    # print(get_line_distance(p0, v0, p1, v1))
